package com.haima.sage.bigdata.api.common;

import java.io.ByteArrayOutputStream;
import java.security.Key;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.Signature;
import java.security.interfaces.RSAKey;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import javax.crypto.Cipher;

public class RSAUtils {
    public static final String KEY_ALGORITHM = "RSA";
    public static final String SIGNATURE_ALGORITHM = "MD5withRSA";
    public static final String PUBLIC_KEY = "RSAPublicKey";
    public static final String PRIVATE_KEY = "RSAPrivateKey";

    public static String encode(Key Key) {
        return encryptBASE64(Key.getEncoded());
    }

    public static byte[] decryptBASE64(String data) {
        return Base64.getDecoder().decode(data);
    }

    public static String encryptBASE64(byte[] data) {
        return new String(Base64.getEncoder().encode(data));
    }

    public static RSAPublicKey getPublicKey(String key) throws NoSuchAlgorithmException, InvalidKeySpecException {
        byte[] keyBytes = decryptBASE64(key);
        X509EncodedKeySpec keySpec = new X509EncodedKeySpec(keyBytes);
        KeyFactory keyFactory = KeyFactory.getInstance("RSA");
        return (RSAPublicKey) keyFactory.generatePublic(keySpec);
    }

    public static RSAPrivateKey getPrivateKey(String key) throws NoSuchAlgorithmException, InvalidKeySpecException {
        byte[] keyBytes = decryptBASE64(key);
        PKCS8EncodedKeySpec pkcs8KeySpec = new PKCS8EncodedKeySpec(keyBytes);
        KeyFactory keyFactory = KeyFactory.getInstance("RSA");
        return (RSAPrivateKey) keyFactory.generatePrivate(pkcs8KeySpec);
    }

    public static String sign(byte[] data, PrivateKey privateKey) throws Exception {
        Signature signature = Signature.getInstance("MD5withRSA");
        signature.initSign(privateKey);
        signature.update(data);
        return encryptBASE64(signature.sign());
    }

    public static boolean verify(byte[] data, PublicKey pubKey, String sign) throws Exception {
        Signature signature = Signature.getInstance("MD5withRSA");
        signature.initVerify(pubKey);
        signature.update(data);

        return signature.verify(decryptBASE64(sign));
    }

    public static byte[] decrypt(byte[] data, RSAPublicKey key) throws Exception {
        int modulus = key.getModulus().bitLength() / 8;
        return decrypt(data, key, modulus);
    }

    public static byte[] decrypt(byte[] data, RSAPrivateKey key) throws Exception {
        int modulus = key.getModulus().bitLength() / 8;
        return decrypt(data, key, modulus);
    }

    public static byte[] decrypt(byte[] data, Key privateKey, int block) {
        try {
            Cipher cipher = Cipher.getInstance("RSA/ECB/PKCS1Padding");
            cipher.init(2, privateKey);

            return doFinal(data, cipher, block);
        } catch (Exception e) {
            e.printStackTrace();
        }
        return null;
    }

    public static String decrypt(String data, RSAPublicKey key) throws Exception {
        byte[] bytes = decryptBASE64(data);
        bytes = decrypt(bytes, key);
        return new String(bytes);
    }

    public static String decrypt(String data, RSAPrivateKey key) throws Exception {
        byte[] bytes = decryptBASE64(data);
        bytes = decrypt(bytes, key);
        return new String(bytes);
    }

    public static byte[] encrypt(byte[] data, RSAPublicKey key) throws Exception {
        int modulus = key.getModulus().bitLength() / 8 - 11;
        return encrypt(data, key, modulus);
    }

    public static byte[] encrypt(byte[] data, RSAPrivateKey key) throws Exception {
        int modulus = key.getModulus().bitLength() / 8 - 11;
        return encrypt(data, key, modulus);
    }

    public static byte[] encrypt(byte[] data, Key publicKey, int block) {
        try {
            Cipher cipher = Cipher.getInstance("RSA/ECB/PKCS1Padding");
            cipher.init(1, publicKey);
            return doFinal(data, cipher, block);
        } catch (Exception e) {
            e.printStackTrace();
        }
        return null;
    }

    public static String encrypt(String data, RSAPublicKey key) throws Exception {
        byte[] bytes = encrypt(data.getBytes(), key);
        return encryptBASE64(bytes);
    }

    public static String encrypt(String data, RSAPrivateKey key) throws Exception {
        byte[] bytes = encrypt(data.getBytes(), key);
        return encryptBASE64(bytes);
    }

    private static byte[] doFinal(byte[] data, Cipher cipher, int block) {
        try {
            int inputLen = data.length;
            ByteArrayOutputStream out = new ByteArrayOutputStream();
            int offSet = 0;

            int i = 0;

            while (inputLen - offSet > 0) {
                byte[] cache;
                if (inputLen - offSet > block)
                    cache = cipher.doFinal(data, offSet, block);
                else {
                    cache = cipher.doFinal(data, offSet, inputLen - offSet);
                }
                out.write(cache, 0, cache.length);
                i++;
                offSet = i * block;
            }
            byte[] distData = out.toByteArray();
            out.close();
            return distData;
        } catch (Exception e) {
            e.printStackTrace();
        }
        return null;
    }

    public static Map<String, String> initKey() throws Exception {
        KeyPairGenerator keyPairGen = KeyPairGenerator.getInstance("RSA");
        keyPairGen.initialize(2048);

        KeyPair keyPair = keyPairGen.generateKeyPair();

        RSAPublicKey rsaPublicKey = (RSAPublicKey) keyPair.getPublic();
        RSAPrivateKey rsaPrivateKey = (RSAPrivateKey) keyPair.getPrivate();

        String publicKey = new String(Base64.getEncoder().encode(rsaPublicKey.getEncoded()));
        String privateKey = new String(Base64.getEncoder().encode(rsaPrivateKey.getEncoded()));

        Map<String, String> keyMap = new HashMap<>(2);
        keyMap.put("RSAPublicKey", publicKey);
        keyMap.put("RSAPrivateKey", privateKey);
        return keyMap;
    }
}